#!/bin/bash

# HPA (Horizontal Pod Autoscaler) 功能测试脚本
# 用法: ./test-hpa.sh [namespace] [service_name] [service_port]

set -e

# 默认参数
NAMESPACE="${1:-nju39}"
SERVICE_NAME="${2:-prometheus-test-demo}"
SERVICE_PORT="${3:-8080}"
HPA_NAME="prometheus-test-hpa"

NODE_IP="172.22.83.19"
NODE_PORT="31073"
SERVICE_URL="http://$NODE_IP:$NODE_PORT"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查必要的工具
check_prerequisites() {
    log_info "检查必要工具..."
    
    if ! command -v kubectl &> /dev/null; then
        log_error "kubectl 未安装或不在PATH中"
        exit 1
    fi
    
    if ! command -v curl &> /dev/null; then
        log_error "curl 未安装或不在PATH中"
        exit 1
    fi
    
    log_success "必要工具检查完成"
}

# 获取服务URL
get_service_url() {
    # log_info "获取服务访问地址..."
    
    # # 尝试获取NodePort服务
    # NODE_PORT=$(kubectl get svc $SERVICE_NAME -n $NAMESPACE -o jsonpath='{.spec.ports[0].nodePort}' 2>/dev/null || echo "")
    
    # if [ -n "$NODE_PORT" ]; then
    #     # 获取节点IP
    #     NODE_IP=$(kubectl get nodes -o jsonpath='{.items[0].status.addresses[?(@.type=="InternalIP")].address}')
    #     SERVICE_URL="http://$NODE_IP:$NODE_PORT"
    # else
    #     # 尝试获取ClusterIP
    #     CLUSTER_IP=$(kubectl get svc $SERVICE_NAME -n $NAMESPACE -o jsonpath='{.spec.clusterIP}' 2>/dev/null || echo "")
    #     if [ -n "$CLUSTER_IP" ]; then
    #         SERVICE_URL="http://$CLUSTER_IP:$SERVICE_PORT"
    #     else
    #         log_error "无法获取服务访问地址"
    #         exit 1
    #     fi
    # fi
    
    log_success "服务URL: $SERVICE_URL"
}

# 检查HPA状态
check_hpa_status() {
    log_info "检查HPA状态..."
    
    if ! kubectl get hpa $HPA_NAME -n $NAMESPACE &> /dev/null; then
        log_error "HPA '$HPA_NAME' 在命名空间 '$NAMESPACE' 中不存在"
        exit 1
    fi
    
    echo "当前HPA状态:"
    kubectl get hpa $HPA_NAME -n $NAMESPACE
    echo ""
    
    log_success "HPA状态检查完成"
}

# 检查初始Pod数量
check_initial_pods() {
    log_info "检查初始Pod状态..."
    
    INITIAL_POD_COUNT=$(kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE -o wide --no-headers | wc -l)
    echo "当前Pod数量: $INITIAL_POD_COUNT"
    
    kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE
    echo ""
    
    log_success "初始状态检查完成"
}

# 检查服务健康状态
check_service_health() {
    log_info "检查服务健康状态..."
    
    for i in {1..5}; do
        if curl -s "$SERVICE_URL/test/health" &> /dev/null; then
            log_success "服务健康检查通过"
            return 0
        fi
        log_warning "健康检查失败，重试 $i/5..."
        sleep 2
    done
    
    log_error "服务健康检查失败"
    exit 1
}

# 监控HPA状态变化
monitor_hpa() {
    local duration=$1
    local interval=${2:-5}
    
    log_info "监控HPA状态变化 (${duration}秒，每${interval}秒检查一次)..."
    
    local end_time=$((SECONDS + duration))
    
    while [ $SECONDS -lt $end_time ]; do
        local current_pods=$(kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE --no-headers | wc -l)
        local hpa_status=$(kubectl get hpa $HPA_NAME -n $NAMESPACE --no-headers)
        
        echo "[$(date +'%H:%M:%S')] Pods: $current_pods | HPA: $hpa_status"
        sleep $interval
    done
}

# 轻负载测试
light_load_test() {
    log_info "执行轻负载测试..."
    
    local duration=60
    local concurrent=5
    local total_requests=100
    
    echo "测试参数: 并发=$concurrent, 总请求=$total_requests, 持续时间=${duration}秒"
    
    # 启动负载测试
    for i in $(seq 1 $concurrent); do
        (
            local count=$((total_requests / concurrent))
            for j in $(seq 1 $count); do
                curl -s "$SERVICE_URL/hello" > /dev/null 2>&1
                sleep 0.5
            done
        ) &
    done
    
    # 监控期间的HPA状态
    monitor_hpa $duration 10 &
    local monitor_pid=$!
    
    wait # 等待所有负载测试完成
    kill $monitor_pid 2>/dev/null || true
    
    log_success "轻负载测试完成"
}

# 高负载测试 - 触发扩容
high_load_test() {
    log_info "执行高负载测试以触发扩容..."
    
    local duration=180
    local concurrent=30
    
    echo "测试参数: 并发=$concurrent, 持续时间=${duration}秒"
    echo "目标: 触发CPU/QPS扩容条件"
    
    # 记录开始时间和初始Pod数量
    local start_pods=$(kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE --no-headers | wc -l)
    echo "扩容前Pod数量: $start_pods"
    
    # 启动高并发负载测试
    for i in $(seq 1 $concurrent); do
        (
            local end_time=$((SECONDS + duration))
            while [ $SECONDS -lt $end_time ]; do
                curl -s "$SERVICE_URL/hello" > /dev/null 2>&1 &
                curl -s "$SERVICE_URL/diagnostic" > /dev/null 2>&1 &
                curl -s "$SERVICE_URL/test/health" > /dev/null 2>&1 &
                sleep 0.1
            done
            wait
        ) &
    done
    
    # 监控扩容过程
    monitor_hpa $duration 15 &
    local monitor_pid=$!
    
    wait # 等待所有负载测试完成
    kill $monitor_pid 2>/dev/null || true
    
    # 检查扩容结果
    sleep 30 # 等待扩容稳定
    local end_pods=$(kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE --no-headers | wc -l)
    echo ""
    echo "扩容后Pod数量: $end_pods"
    
    if [ $end_pods -gt $start_pods ]; then
        log_success "✅ 扩容成功! Pod数量从 $start_pods 增加到 $end_pods"
    else
        log_warning "⚠️  未观察到扩容，可能需要更高负载或更长时间"
    fi
}

# 缩容测试
scale_down_test() {
    log_info "执行缩容测试..."
    
    local current_pods=$(kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE --no-headers | wc -l)
    
    if [ $current_pods -le 1 ]; then
        log_warning "当前只有 $current_pods 个Pod，跳过缩容测试"
        return
    fi
    
    echo "当前Pod数量: $current_pods"
    echo "停止负载，等待自动缩容..."
    echo "注意: 缩容有5分钟的稳定窗口期"
    
    # 监控缩容过程 (10分钟)
    monitor_hpa 600 30 &
    local monitor_pid=$!
    
    # 等待一段时间后检查结果
    sleep 600
    kill $monitor_pid 2>/dev/null || true
    
    local final_pods=$(kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE --no-headers | wc -l)
    echo ""
    echo "缩容后Pod数量: $final_pods"
    
    if [ $final_pods -lt $current_pods ]; then
        log_success "✅ 缩容成功! Pod数量从 $current_pods 减少到 $final_pods"
    else
        log_warning "⚠️  未观察到缩容，可能需要更长等待时间"
    fi
}

# 生成测试报告
generate_report() {
    log_info "生成测试报告..."
    
    local report_file="hpa_test_report/hpa_test_report_$(date +%Y%m%d_%H%M%S).txt"
    
    {
        echo "========================================"
        echo "HPA 功能测试报告"
        echo "========================================"
        echo "测试时间: $(date)"
        echo "命名空间: $NAMESPACE"
        echo "服务名称: $SERVICE_NAME"
        echo "HPA名称: $HPA_NAME"
        echo "服务URL: $SERVICE_URL"
        echo ""
        
        echo "最终HPA状态:"
        kubectl get hpa $HPA_NAME -n $NAMESPACE
        echo ""
        
        echo "最终Pod状态:"
        kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE
        echo ""
        
        echo "HPA详细信息:"
        kubectl describe hpa $HPA_NAME -n $NAMESPACE
        echo ""
        
        echo "Pod资源使用情况:"
        kubectl top pods -l app=$SERVICE_NAME -n $NAMESPACE 2>/dev/null || echo "metrics-server 不可用"
        
    } > $report_file
    
    log_success "测试报告已保存到: $report_file"
}

# 主函数
main() {
    echo "========================================"
    echo "🚀 HPA 功能测试脚本"
    echo "========================================"
    echo "命名空间: $NAMESPACE"
    echo "服务名称: $SERVICE_NAME"
    echo "HPA名称: $HPA_NAME"
    echo ""
    
    # 执行测试步骤
    check_prerequisites
    get_service_url
    check_hpa_status
    check_initial_pods
    check_service_health
    
    echo ""
    echo "========================================"
    echo "📊 开始负载测试"
    echo "========================================"
    
    # 轻负载测试
    light_load_test
    sleep 30
    
    # 高负载测试 - 触发扩容
    high_load_test
    sleep 60
    
    # 缩容测试
    echo ""
    read -p "是否执行缩容测试？(需要等待约10分钟) [y/N]: " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        scale_down_test
    else
        log_info "跳过缩容测试"
    fi
    
    # 生成报告
    generate_report
    
    echo ""
    echo "========================================"
    echo "🎉 HPA 测试完成!"
    echo "========================================"
    
    # 最终状态显示
    echo "最终状态:"
    kubectl get hpa $HPA_NAME -n $NAMESPACE
    kubectl get pods -l app=$SERVICE_NAME -n $NAMESPACE
}

# 信号处理 - 清理后台进程
trap 'jobs -p | xargs -r kill; exit 130' INT TERM

# 执行主函数
main "$@" 